from math import sqrt
import numpy as np
from PQRSTCheck import PQRST


# RR心拍划分
def splitByR(ecg, checkR):
    index = checkR
    ecg_rhythm = []
    for i in range(1, len(index) - 1):
        # index[i-1] 和 index[i+1]：分别代表当前索引位置 i 的前一个和后一个 R 波的位置。
        ecg_rhythm.append(ecg[index[i - 1]:index[i + 1]])
    for i in range(len(ecg_rhythm)):
        if len(ecg_rhythm[i]) > 470:
            ecg_rhythm[i] = ecg_rhythm[i][:470]
        elif len(ecg_rhythm[i]) < 470:
            average = np.mean(ecg_rhythm[i])
            # 用于创建一个具有特定形状且所有元素都填充为指定值的新数组，average 用来填充新数组的值。
            average_arr = np.full(470 - len(ecg_rhythm[i]), average)
            ecg_rhythm[i] = np.concatenate((ecg_rhythm[i], average_arr))
    return np.array(ecg_rhythm)


# 提取距离，幅值特征
def feature(ecgdata, locatedR, locatedQ, locatedS, locatedP, locatedT, locatedPBegin, locatedPEnd, locatedTBegin,
            locatedTEnd):
    featureVector = []  # 特征向量，用于存储间期特征
    amplitudeVector = []  # 特征向量，用于存储幅值特征
    areaVector = []  # 特征向量，用于存储归一化后的面积特征
    for i in range(len(locatedR) - 1):
        # =========================面积计算
        # P波段
        try:
            p_segment = 0
            if len(locatedPBegin) > i:
                p_segment = ecgdata[locatedPBegin[i]:locatedPEnd[i]]
            p_baseline = np.mean(p_segment)
            p_corrected = p_segment - p_baseline
            # 使用np.trapz进行梯形积分（比矩形法更精确）
            if np.size(p_corrected) == 0:
                print("警告: p_corrected 是空数组，无法计算积分。跳过或使用默认值。")
                areaP = 0.0  # 或其它合适的默认值
            else:
                areaP = np.trapz(p_corrected)
        except Exception as e:
            areaP = 0.0
        # QRS段
        qrs_segment = ecgdata[locatedQ[i]:locatedS[i]]
        qrs_baseline = np.mean(qrs_segment)
        qrs_corrected = qrs_segment - qrs_baseline
        areaQRS = np.trapz(qrs_corrected)
        # T波段
        t_segment = ecgdata[locatedTBegin[i]:locatedTEnd[i]]
        t_baseline = np.mean(t_segment)
        t_corrected = t_segment - t_baseline
        areaT = np.trapz(t_corrected)
        areaVector.append([areaP, areaQRS, areaT])

        # 计算各个时点的心电图特征值
        # abs() 取绝对值
        # 前一个 RR 间期，反映前次心跳的周期长度，是计算心率变异性 (HRV) 的基础。
        pre_RR = abs(locatedR[i] - locatedR[i - 1])
        # 后一个 RR 间期，反映下次心跳的周期长度，同样用于 HRV 分析
        post_RR = abs(locatedR[i + 1] - locatedR[i])
        PQ = abs(locatedQ[i] - locatedP[i])
        # 房室传导时间。从心房除极开始到心室除极开始的时间，延长可能提示房室传导阻滞。
        PR = abs(locatedR[i] - locatedP[i])
        PS = abs(locatedS[i] - locatedP[i])
        PT = abs(locatedT[i] - locatedP[i])
        QR = abs(locatedR[i] - locatedQ[i])
        QS = abs(locatedS[i] - locatedQ[i])
        QT = abs(locatedQ[i] - locatedT[i])
        QT = QT / sqrt((pre_RR + post_RR) / 2)
        RS = abs(locatedS[i] - locatedR[i])
        RT = abs(locatedT[i] - locatedR[i])
        ST = abs(locatedT[i] - locatedS[i])
        PQR = abs(locatedR[i] - locatedP[i])
        QRS = abs(locatedS[i] - locatedQ[i])
        RST = abs(locatedT[i] - locatedR[i])
        PQRS = abs(locatedS[i] - locatedP[i])
        QRST = abs(locatedT[i] - locatedQ[i])
        _PP_ = abs(locatedPEnd[i] - locatedPBegin[i])
        _TT_ = abs(locatedTEnd[i] - locatedTBegin[i])
        _PQ = abs(locatedQ[i] - locatedPBegin[i])
        ST_ = abs(locatedTEnd[i] - locatedS[i])
        _PR = abs(locatedR[i] - locatedPBegin[i])
        P_R = abs(locatedR[i] - locatedPEnd[i])
        R_T = abs(locatedTBegin[i] - locatedR[i])
        RT_ = abs(locatedTEnd[i] - locatedR[i])
        # featureVector.append([pre_RR,post_RR])
        featureVector.append(
            [pre_RR, post_RR, PQ, PR, PS, PT, QR, QS, QT, RS, RT, ST, PQR, QRS, RST, PQRS, QRST, _PP_, _TT_, _PQ, ST_,
             _PR, P_R, R_T, RT_])

        # 计算各个时点的心电图幅值特征
        ampRR = ecgdata[locatedR[i + 1]] - ecgdata[locatedR[i]]
        ampPQ = ecgdata[locatedQ[i]] - ecgdata[locatedP[i]]
        ampPR = ecgdata[locatedR[i]] - ecgdata[locatedP[i]]
        ampPS = ecgdata[locatedS[i]] - ecgdata[locatedP[i]]
        ampPT = ecgdata[locatedT[i]] - ecgdata[locatedP[i]]
        ampQR = ecgdata[locatedR[i]] - ecgdata[locatedQ[i]]
        ampQS = ecgdata[locatedS[i]] - ecgdata[locatedQ[i]]
        ampQT = ecgdata[locatedT[i]] - ecgdata[locatedQ[i]]
        ampRS = ecgdata[locatedS[i]] - ecgdata[locatedR[i]]
        ampRT = ecgdata[locatedT[i]] - ecgdata[locatedR[i]]
        ampST = ecgdata[locatedT[i]] - ecgdata[locatedS[i]]
        ampPQR = ecgdata[locatedR[i]] - ecgdata[locatedP[i]]
        ampQRS = ecgdata[locatedS[i]] - ecgdata[locatedQ[i]]
        ampRST = ecgdata[locatedT[i]] - ecgdata[locatedR[i]]
        ampPQRS = ecgdata[locatedS[i]] - ecgdata[locatedP[i]]
        ampQRST = ecgdata[locatedT[i]] - ecgdata[locatedQ[i]]
        amp_PP = ecgdata[locatedP[i]] - ecgdata[locatedPBegin[i]]
        ampPP_ = ecgdata[locatedPEnd[i]] - ecgdata[locatedP[i]]
        amp_TT = ecgdata[locatedT[i]] - ecgdata[locatedTBegin[i]]
        ampTT_ = ecgdata[locatedTEnd[i]] - ecgdata[locatedT[i]]
        amp_PR = ecgdata[locatedR[i]] - ecgdata[locatedPBegin[i]]
        ampRT_ = ecgdata[locatedTEnd[i]] - ecgdata[locatedR[i]]
        amp_PQ = ecgdata[locatedQ[i]] - ecgdata[locatedPBegin[i]]
        ampST_ = ecgdata[locatedTEnd[i]] - ecgdata[locatedS[i]]

        amplitudeVector.append(
            [ampRR, ampPQ, ampPR, ampPS, ampPT, ampQR, ampQS, ampQT, ampRS, ampRT, ampST, ampPQR, ampQRS, ampRST,
             ampPQRS, ampQRST, amp_PP, ampPP_, amp_TT, ampTT_, amp_PR, ampRT_, amp_PQ, ampST_])

    # 归一化面积特征向量
    areaVector = np.array(areaVector)
    maxFeature = np.max(areaVector)
    minFeature = np.min(areaVector)
    areaVector = 2 * (areaVector - minFeature) / (maxFeature - minFeature) - 1

    # 归一化特征向量
    featureVector = np.array(featureVector)
    maxFeature = np.max(featureVector)
    minFeature = np.min(featureVector)
    featureVector = 2 * (featureVector - minFeature) / (maxFeature - minFeature) - 1

    # 归一化幅值特征向量
    amplitudeVector = np.array(amplitudeVector)
    maxAmplitude = np.max(amplitudeVector)
    minAmplitude = np.min(amplitudeVector)
    amplitudeVector = 2 * (amplitudeVector - minAmplitude) / (maxAmplitude - minAmplitude) - 1

    return areaVector, featureVector, amplitudeVector


def main_feature(ecg, locatedR, sampling_rate):
    # RR心拍划分, RR_interval是二维数组，代表每个心拍的470个点数据，不足的部分由平均值补足
    RR_interval = splitByR(ecg, locatedR)

    #  PQRST波检测  locatedR:2691
    #  2689 舍弃了第一个波形和最后一个波形
    #  2691      2691      2689       2689     2689            2689        2689               2689
    _, locatedQ, locatedS, locatedP, locatedT, locatedPBegin, locatedPEnd, locatedTBegin, locatedTEnd = PQRST(ecg,
                                                                                                              locatedR,
                                                                                                              sampling_rate)
    # 面积特征（3个）， 间期特征（25个），幅值特征（24个）
    # areaVector       面积特征：归一化后的面积特征，包括 P波，QRS波，T波面积
    # featureVector    间期特征  RR间期前半部分面积和后半部分面积
    # amplitudeVector  幅值特征 ，幅值是在一个周期内最大绝对值
    # ampRR, ampPQ, ampPR, ampPS, ampPT, ampQR, ampQS, ampQT, ampRS,
    # ampRT, ampST, ampPQR, ampQRS, ampRST,ampPQRS, ampQRST, amp_PP, ampPP_, amp_TT, ampTT_,
    # amp_PR, ampRT_, amp_PQ, ampST
    areaVector, featureVector, amplitudeVector = feature(ecg, locatedR[1:], locatedQ[1:], locatedS[1:], locatedP,
                                                         locatedT,
                                                         locatedPBegin, locatedPEnd, locatedTBegin, locatedTEnd)

    return areaVector, featureVector, amplitudeVector, RR_interval
